package com.xz.xzmybatisagent;

import com.xz.xzmybatisagent.utils.ClassScanConfig;
import com.xz.xzmybatisagent.utils.ClassScanner;
import javassist.*;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author DONG.WANG17@HAND-CHINA.COM
 * @since 2025-02-13 17:01
 */
public class InterceptorClassFileTransformer implements ClassFileTransformer {
    private int sqlPort;
    private int activePort;
    // 排除注入的类
    private final List<String> injectExcludes = new ArrayList<>();
    private final ClassScanConfig scanConfig = new ClassScanConfig();

    {
        injectExcludes.add("com.xz.xzmybatisagent.InterceptorClassFileTransformer");
        injectExcludes.add("com.xz.xzmybatisagent.AgentmainAgent");
        injectExcludes.add("com.xz.xzmybatisagent.PremainAgent");
        injectExcludes.add("com.xz.xzmybatisagent.socket.example.ClientExample");
        injectExcludes.add("com.xz.xzmybatisagent.socket.example.ServerExample");

        scanConfig.setPackagePrefix("com.xz.xzmybatisagent");
        scanConfig.setExcludedClasses(injectExcludes);
        scanConfig.setExcludedDirectories(Collections.singletonList("xzmybatisagent/utils/"));
        scanConfig.setIncludeInnerClasses(true);
    }

    public void setArgs(String agentArgs) {
        // sqlPort=55389,activePort=55486
        System.out.println("Sql Parser started at port\t" + agentArgs);
        String[] args = agentArgs.split(",");
        for (String arg : args) {
            if (arg.startsWith("sqlPort=")) {
                sqlPort = Integer.parseInt(arg.trim().replace("sqlPort=", ""));
            }
            if (arg.startsWith("activePort=")) {
                activePort = Integer.parseInt(arg.trim().replace("activePort=", ""));
            }
        }
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        if ("org/apache/ibatis/plugin/InterceptorChain".equals(className)) {
            System.out.println("[JavaAgent] MyBatis 拦截器注入中...");
            CtClass ctClass = null;
            try {
                // 强制使用 loader 作为 classloader 解决类加载器隔离导致找不到类的问题
                final List<String> classes = ClassScanner.getClassNameFromJar(ClassScanner.class.getProtectionDomain().getCodeSource().getLocation(), scanConfig);
                injectClassForName(loader, classes);

                // ClassPool classPool = ClassPool.getDefault();
                // 独立类池，避免缓存
                ClassPool classPool = new ClassPool(true);

                classPool.appendClassPath(new ClassClassPath(this.getClass()));
                classPool.appendClassPath(new LoaderClassPath(this.getClass().getClassLoader()));
                classPool.appendClassPath(new ClassClassPath(EnhanceUtils.class));
                classPool.appendClassPath(new LoaderClassPath(EnhanceUtils.class.getClassLoader()));
                // 使用目标类加载器
                classPool.insertClassPath(new LoaderClassPath(loader));

                ctClass = classPool.get("org.apache.ibatis.plugin.InterceptorChain");

                // 2、插入自定义逻辑
                final CtMethod pluginAll = ctClass.getDeclaredMethod("pluginAll");
                pluginAll.insertBefore("com.xz.xzmybatisagent.EnhanceUtils.addMyInterceptor($0.interceptors," + sqlPort + "," + activePort + ");");

                System.out.println("[JavaAgent] MyBatis 拦截器注入成功");

                return ctClass.toBytecode();
            } catch (Throwable t) {
                System.err.println("[JavaAgent] MyBatis 拦截器注入失败: " + t.getMessage());
            } finally {
                if (ctClass != null) {
                    ctClass.detach();
                }
            }
        }
        return classfileBuffer;
    }

    private void injectClassForName(ClassLoader loader, Iterable<String> classNames) {
        for (String name : classNames) {
            try {
                // 先尝试直接加载
                loader.loadClass(name);
            } catch (Throwable e) {
                try {
                    Class<?> targetClass = Class.forName(name);
                    // 递归注入父类和接口
                    injectSuperClassesAndInterfaces(loader, targetClass);
                    // 最后注入目标类本身
                    injectClassToLoader(loader, targetClass);
                } catch (Throwable ex) {
                    System.err.printf("Failed to load class [%s] with error message: %s%n", name, ex.getMessage());
                }
            }
        }
    }

    private void injectSuperClassesAndInterfaces(ClassLoader loader, Class<?> clazz) {
        // 处理父类
        Class<?> superClass = clazz.getSuperclass();
        if (superClass != null && !superClass.equals(Object.class)) {
            try {
                if (loader.loadClass(superClass.getName()) == null) {
                    injectSuperClassesAndInterfaces(loader, superClass);
                    injectClassToLoader(loader, superClass);
                }
            } catch (ClassNotFoundException e) {
                injectSuperClassesAndInterfaces(loader, superClass);
                injectClassToLoader(loader, superClass);
            }
        }

        // 处理接口
        for (Class<?> interfaceClass : clazz.getInterfaces()) {
            try {
                if (loader.loadClass(interfaceClass.getName()) == null) {
                    injectSuperClassesAndInterfaces(loader, interfaceClass);
                    injectClassToLoader(loader, interfaceClass);
                }
            } catch (ClassNotFoundException e) {
                injectSuperClassesAndInterfaces(loader, interfaceClass);
                injectClassToLoader(loader, interfaceClass);
            }
        }
    }

    /**
     * 动态注入类
     *
     * @param loader 注入的加载器
     * @param clazz  需要注入的类
     */
    private void injectClassToLoader(ClassLoader loader, Class<?> clazz) {
        try {
            String className = clazz.getName().replace('.', '/');
            ByteArrayOutputStream buffer = null;
            try (InputStream is = clazz.getClassLoader().getResourceAsStream(className + ".class")) {
                if (is == null) return;

                buffer = new ByteArrayOutputStream();
                int nRead;
                byte[] data = new byte[1024];
                while ((nRead = is.read(data, 0, data.length)) != -1) {
                    buffer.write(data, 0, nRead);
                }
            } catch (Throwable t) {
                System.err.printf("Error reading class file: %s", t.getMessage());
            }
            buffer.flush();
            byte[] bytes = buffer.toByteArray();

            Method defineClass = ClassLoader.class.getDeclaredMethod("defineClass", String.class, byte[].class, int.class, int.class);
            defineClass.setAccessible(true);
            defineClass.invoke(loader, clazz.getName(), bytes, 0, bytes.length);

            System.out.printf("[%s]====================== 动态注入成功%n", clazz.getName());
        } catch (Throwable t) {
            t.printStackTrace();
            System.err.printf("[%s]====================== 动态注入失败: %s%n", clazz.getName(), t.getMessage());
        }
    }
}
